
import sys
from config import *
import coveragefromsummary

MAX_GAP = 100

def lookupGeneInterval(gene, filename=KNOWN_GENE_TABLE):
    for line in open(filename):
        if line[0]=="#": continue
        cols = line.strip().split("\t")
        if gene != cols[0]: continue
        return (cols[1], cols[2], long(cols[3]), long(cols[4]))
    print >> sys.stderr, "# Couldn't find gene ", gene
    return None

def trackName(queryId, samplefile, samplename):
    if samplefile != samplename:
        return queryId + "--" + samplename
    else:
        return queryId + "--" +samplefile.split("/")[-1]
    
def generateWiggleCoverage(requests, samplefiles, filename=KNOWN_GENE_TABLE):

    queries = {}
    for q in requests:
        if len(q)==3:
            # query of form ("chr21", 1, 1000)
            (chr, spos, epos) = (q[0], long(q[1]), long(q[2]))
            queryId = "%s:%s-%s" % (chr, spos, epos)
            queries[queryId] = (chr, '?', spos, epos)
        elif len(q)==2:
            (gene, utrLength) = (q[0], long(q[1]))
            (chr, strand, spos, epos) = lookupGeneInterval(gene)            
            queries[gene] = (chr, '?', spos - utrLength, epos + utrLength)
        else:
            raise Exception("bad coverage request: " + str(q))
    
    for queryId in queries.keys():
        print >> sys.stderr, "Query '%s' @ %s" % (queryId, queries[queryId])

    printedHeader = False
    printedInitialStep = False
    for queryId in queries.keys():
        print >> sys.stderr, "Processing query '%s' @ %s ..." % (queryId, queries[queryId])
        (chr, strand, spos, epos) = queries[queryId]
        for (samplefile, samplename) in samplefiles:
            if not printedHeader:
                printedHeader = True
                print "track type=wiggle_0 name=\"%s\" description=\"%s\"" % (queryId, trackName(queryId, samplefile, samplename))
            coverage = coveragefromsummary.process(queries[queryId], samplefile)
            unprintedZeroes = 0
            for (x, cov) in coverage:
                if cov.total == 0:
                    unprintedZeroes += 1
                    continue
                if unprintedZeroes > MAX_GAP or not printedInitialStep:
                    print "fixedStep chrom=%s start=%ld step=1" % (chr, x)
                    printedInitialStep = True
                    unprintedZeroes = 0
                while unprintedZeroes > 0:
                    print "0"
                    unprintedZeroes -= 1
                print cov
            if not printedInitialStep:
                print "fixedStep chrom=%s start=%ld step=1" % (chr, spos)
                print "0"

def printUsage(args, usage, minArgs=2):
  if len(args) < minArgs or args[1].startswith("-h") or args[1].startswith("--h"):
    print usage
    sys.exit(1)

usageString = """
Usage: python %s QUERY SUMMARY_FILE1 SUMMARY_FILE2 ...,
where QUERY is either [-UTR utrLength] GENE or -Q CHR MIN_POS MAX_POS.

Generates a coverage map for the specified gene or interval, in wiggle format,
from each of the specified summary files.  (If -UTR is specified, extend
the gene on either side by the given UTR length.)
"""

def getSampleName(samplefile):
    filename = samplefile.split("/")[-1]
    ignoreBits = ["illumina", "burge", "txt", "snp-results", "bowtie_sorted_summary", "INDEX"]
    bits = filter(lambda(x):(x not in ignoreBits), filename.split("."))
    return ".".join(bits)

if __name__=="__main__":
    printUsage(sys.argv, usageString % sys.argv[0], 3)
    if sys.argv[1]=='-Q':
        chr = sys.argv[2]
        spos = long(sys.argv[3])
        epos = long(sys.argv[4])
        query = (chr, spos, epos)
        samples = sys.argv[5:]
    elif sys.argv[1]=='-UTR':
        utrLength = int(sys.argv[2])
        gene = sys.argv[3]
        query = (gene, utrLength)
        samples = sys.argv[4:]
    else:
        gene = sys.argv[1]
        query = (gene, 0)
        samples = sys.argv[2:]
    generateWiggleCoverage([query], [(x,getSampleName(x)) for x in samples])
